from paddlex_table.text_detection import text_detection
import numpy as np
import fitz  # PyMuPDF
from wired_table_rec.utils.utils import VisTable
from table_cls import TableCls
from wired_table_rec.main import WiredTableInput, WiredTableRecognition
from lineless_table_rec.main import LinelessTableInput, LinelessTableRecognition
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
import logging  
# 初始化组件（只执行一次）
wired_input = WiredTableInput(device="cuda")
lineless_input = LinelessTableInput(device="cuda")
wired_engine = WiredTableRecognition(wired_input)
lineless_engine = LinelessTableRecognition(lineless_input)
table_cls = TableCls(model_type='paddle')

viser = VisTable()

def get_text_from_rect(page, page_size, poly):
    """
    按指定多边形区域，从PDF页面提取文本，带防错机制。
    Extract text from a given polygon region of a PDF page, with robust error handling.

    Args:
        page (fitz.Page): 
            PDF页面对象，PyMuPDF的page实例。  
            PDF page object, as from PyMuPDF.
        page_size (tuple): 
            页面尺寸，(width, height)，与原始图片一致。  
            Size of the page, (width, height), matching the original image.
        poly (array-like): 
            区域多边形点的列表或数组，格式为[[x1, y1], [x2, y2], ...]。  
            Polygon points, list/array of [[x1, y1], [x2, y2], ...].

    Returns:
        str:
            提取到的文本内容（可能为空字符串）。  
            Extracted text content, or an empty string on error.
    """
    try:
        # 校验参数有效性
        if not hasattr(page, 'rect'):
            raise ValueError("page对象不含rect属性")
        if not isinstance(page_size, (tuple, list)) or len(page_size) != 2:
            raise ValueError(f"page_size格式错误: {page_size}")
        if not hasattr(poly, '__iter__') or len(poly) < 2:
            raise ValueError(f"poly必须是二维点列表，实际: {poly}")

        width_mineru, height_mineru = page_size
        if width_mineru == 0 or height_mineru == 0:
            raise ValueError(f"page_size非法: {page_size}")

        # 提取多边形的x/y坐标分量
        x_coords = [point[0] for point in poly]
        y_coords = [point[1] for point in poly]

        rect = page.rect
        width = rect.width
        height = rect.height
        # 坐标归一化变换，映射回PDF真实坐标
        x0 = min(x_coords) * (width / width_mineru)
        y0 = min(y_coords) * (height / height_mineru)
        x1 = max(x_coords) * (width / width_mineru)
        y1 = max(y_coords) * (height / height_mineru)
        rect = fitz.Rect(x0, y0, x1, y1)

        # 提取指定区域文本
        text = page.get_text("text", clip=rect)
        if text is None:
            text = ""
        return text

    except Exception as e:
        logging.warning(f"警告：区域文本提取失败，poly={poly}, 错误信息: {e}")
        return ""


def process_single_ocr(task, text_coord, doc):
    """
    单个表格的并行处理逻辑  
    Process a single table OCR task in parallel with robust error handling.

    Args:
        task (dict): 
            任务字典，包含表格的坐标、页码等信息。  
            Task dictionary, including table coordinates, page number, etc.
        text_coord (dict): 
            文本坐标字典，包含文本块的坐标、置信分数等信息。  
            Dictionary with text block coordinates and confidence scores.
        doc (fitz.Document): 
            PDF文档对象，由fitz打开。  
            PDF document object, opened by fitz.

    Returns:
        list:
            每个元素为三元组 (dt_poly, text, score)：
                - dt_poly (np.ndarray): 原始图片坐标（多边形点）
                - text (str): 提取的文本
                - score (float): 置信分数  
            List of (dt_poly, text, score) tuples, where:
                - dt_poly (np.ndarray): original polygon coordinates in image
                - text (str): extracted text
                - score (float): confidence score
            出现异常时，返回空列表 []。
    """
    try:
        # 1. 检查必要的字段
        if not all(k in task for k in ('page_no', 'page_size', 'layout_det')):
            raise ValueError(f"任务缺少必要字段: {task}")
        if 'poly' not in task['layout_det']:
            raise ValueError(f"layout_det 字段缺 poly: {task['layout_det']}")
        if not all(k in text_coord for k in ('img_size', 'dt_polys', 'dt_scores')):
            raise ValueError(f"text_coord 缺必要字段: {text_coord}")

        page_no = task['page_no']
        if page_no >= len(doc):
            raise IndexError(f"page_no 超出文档页数: {page_no}, 总页数 {len(doc)}")
        page = doc[page_no]
        page_size = task['page_size']
        poly = task['layout_det']['poly']
        if len(poly) < 6:
            raise ValueError(f"poly 点数量不够: {poly}")

        x_min = poly[0]
        x_max = poly[2]
        y_min = poly[1]
        y_max = poly[5]
        img_width_mineru = x_max - x_min
        img_height_mineru = y_max - y_min

        img_size = text_coord['img_size']
        if img_size[0] == 0 or img_size[1] == 0:
            raise ValueError(f"img_size 不合法: {img_size}")
        scale = np.array([img_width_mineru / img_size[0], img_height_mineru / img_size[1]])
        offset = np.array([x_min, y_min])

        
        dt_polys = text_coord['dt_polys']
        dt_scores = text_coord['dt_scores']

        if len(dt_polys) != len(dt_scores):
            raise ValueError(f"dt_polys 和 dt_scores 数量不一致: {len(dt_polys)} vs {len(dt_scores)}")
        
        # dt_start = time.time()
        
        ocr_result = []
        for dt_poly, dt_score in zip(dt_polys, dt_scores):
            try:
                pdf_poly = dt_poly * scale + offset
                text = get_text_from_rect(page, page_size, pdf_poly)
                text = text.rstrip('\n') if isinstance(text, str) else ""
                ocr_result.append((dt_poly, text, dt_score))
            except Exception as e:
                logging.warning(f"警告：文本区域提取失败, 原因: {e}，dt_poly: {dt_poly}")
                ocr_result.append((dt_poly, "", 0))
        # dt_end = time.time()
        # logging.info(f"单个表格处理耗时: {dt_end - dt_start:.2f}秒")
        return ocr_result

    except Exception as e:
        logging.error(f"致命错误：单个表格处理失败！任务: {task}, text_coord: {text_coord}, 错误信息: {e}")
        # 这里可以选择直接返回空，或者raise，看你想怎么兜底
        return []


def table2html(img_paths,tasks,pdf_path,batch_size=1,model_type="auto"):
    """
    表格图片转为HTML结构的主流程函数

    Args:
        img_paths (list): 包含表格图片路径的列表。
        tasks (list): 每个元素为字典，包含每个表格的任务信息，如坐标等。
        pdf_path (str): 源PDF文件的路径。
        batch_size (int): 批处理大小，默认为1。
        model_type (str): 选择表格结构识别模型类型，可选'auto'、'wired'、'lineless'。

    Returns:
        table_results (list): 返回每张图片表格结构识别后的结果（HTML等）。
    """
    # 步骤1：OCR检测图片中的文本块，获取每个文本块的坐标、置信分数和图片尺寸
    start_time1 = time.time()
    text_coords = text_detection(img_paths, batch_size=batch_size, save=False)
    end_time1 = time.time()
    logging.info(f"步骤1，文本块检测耗时: {end_time1 - start_time1:.2f}秒")
    import os
    cpu_count = os.cpu_count()

    # 步骤2：合并PDF坐标映射、文本提取和OCR结果封装
    ocr_results = []
    doc = fitz.open(pdf_path)
    with ThreadPoolExecutor(max_workers=cpu_count*10) as executor:  # max_workers可调整，建议等于CPU核心数或略多
        # 这里用列表生成式把所有任务和参数配好
        futures = [
            executor.submit(process_single_ocr, task, text_coords[i], doc)
            for i, task in enumerate(tasks)
        ]
        # 保证顺序不乱，按原顺序收集结果
        for future in futures:
            ocr_results.append(future.result()) 
    end_time2 = time.time()
    logging.info(f"步骤2，OCR处理耗时: {end_time2 - end_time1:.2f}秒")


    # 步骤3：对每张图片执行表格结构识别（wired/lineless/auto自动选择），返回最终HTML结构等
    # 开始计时
    start_time3 = time.time()
    table_results = []
    for i, img_path in enumerate(img_paths):
        ocr_result = ocr_results[i]
        # 自动选择模型
        if model_type == "auto":
            cls, _ = table_cls(img_path)
            table_engine = wired_engine if cls == "wired" else lineless_engine
        elif model_type == "wired":
            table_engine = wired_engine
        elif model_type == "lineless":
            table_engine = lineless_engine
        else:
            raise ValueError("model_type must be one of 'wired', 'lineless', 'auto'")
        table_result = table_engine(img_path, ocr_result=ocr_result)
        table_results.append(table_result)
    logging.info(f"步骤3，表格结构推理耗时: {time.time() - start_time3:.2f}秒")
    return table_results
            

if __name__ == "__main__":
    # 测试代码:运行不了
    img_paths = [
        "path/to/image1.png",
        "path/to/image2.png",
        # 添加更多图片路径
    ]
    table2html(img_paths)